# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

import logging

import experimental.prophet as prophet
from experimental.prophet.ops.fill_query import FillQuery
from coin.exchange.base.tick import FixedTickPrice


def basic(symbol_info, feed):
  book = prophet.book(feed, inverse_product=symbol_info.inverse)
  raw_ask0_p = book.ask0.price.named_ref('raw_ask0_p')
  raw_bid0_p = book.bid0.price.named_ref('raw_bid0_p')
  ask0_p = book.estimated_ask0_price.named_ref('ask0_p')
  bid0_p = book.estimated_bid0_price.named_ref('bid0_p')
  mid_p = ((ask0_p + bid0_p) / 2.).named_ref('mid_p')
  windows = ['1s', '2s', '5s', '10s', '15s', '20s', '30s', '1m', '2m', '5m']
  y_mid_p = prophet.time_shift_y(mid_p, windows)
  return [raw_ask0_p, raw_bid0_p, ask0_p, bid0_p, mid_p, *y_mid_p]


def multiple_edge_fillinfo(symbol_info,
                           feed,
                           theo_ask0_p,
                           theo_bid0_p,
                           trigger,
                           latency,
                           order_life,
                           taker_panalty=0.):
  trigger = trigger or True
  book = prophet.book(feed, inverse_product=symbol_info.inverse)
  trade = prophet.create_node('fastfeed.Trade', [feed.trade]).outputs
  fill_query = FillQuery(book,
                         trade,
                         latency,
                         order_life,
                         maker_fee=symbol_info.maker_fee,
                         taker_fee=symbol_info.taker_fee + taker_panalty)

  tick_price = symbol_info.product_info.gen_tick_price()
  assert isinstance(tick_price, FixedTickPrice)
  tick_size = float(tick_price.tick_size)
  assert tick_size > 0.
  logging.info('%s tick_size: %f', symbol_info.nickname, tick_size)

  vars = []
  with prophet.control_if(trigger):
    for edge_bp in range(-22, 32, 2):
      buy_p = prophet.floor((theo_bid0_p * (1. - edge_bp / 10000.)) / tick_size) * tick_size
      sell_p = prophet.ceil((theo_ask0_p * (1. + edge_bp / 10000.)) / tick_size) * tick_size
      buy_fill, buy_taker, buy_fill_price, buy_fee = fill_query.query_buy(buy_p)
      sell_fill, sell_taker, sell_fill_price, sell_fee = fill_query.query_sell(sell_p)

      vars += [
          buy_p.named_ref('buy_p_%d' % edge_bp),
          buy_fill.named_ref('buy_filled_%d' % edge_bp),
          buy_taker.named_ref('buy_taker_%d' % edge_bp),
          buy_fill_price.named_ref('buy_fill_price_%d' % edge_bp),
          buy_fee.named_ref('buy_fee_%d' % edge_bp),
          sell_p.named_ref('sell_p_%d' % edge_bp),
          sell_fill.named_ref('sell_filled_%d' % edge_bp),
          sell_taker.named_ref('sell_taker_%d' % edge_bp),
          sell_fill_price.named_ref('sell_fill_price_%d' % edge_bp),
          sell_fee.named_ref('sell_fee_%d' % edge_bp)
      ]
  return vars
